from time import perf_counter
import os

from news_classification.download_news_dataset import get_latest_news_keynos, get_news_by_url, batch_get_news_by_oss
from news_classification.lda import train, show_lda_topic, infer_doc_topics, get_lda_model, compute_coherence_value, \
    get_topic_word_matrix, get_word_topic_matrix
from news_classification.glda import get_glda_infer_engine, infer_doc_topics_by_glda, read_glda_phi_file, \
    calculate_glda_mapv, read_glda_theta_file
from news_classification.utils import split_file, convert_sparse2dense
from news_classification.convert_dataset2corpus import process_news_jsonl_corpus
from news_classification.dependent_scripts.convert_gensim_lda2familia import convert_gensim_lda2familia


def test_get_latest_news_keynos():
    start = perf_counter()
    get_latest_news_keynos()
    print('Consume: %f' % (perf_counter() - start))


def test_get_news_by_url():
    keyno = 'a9ac2738a329ffb7423f9678e267d4b6'
    res = get_news_by_url(keyno)
    print(res)


def test_batch_get_news_by_oss():
    keynos = get_latest_news_keynos(10000)
    batch_get_news_by_oss(keynos, output='data/1w-data.jsonl')


def test_prepare_corpus():
    process_news_jsonl_corpus('data/5M-data.jsonl', 'data/5M-data-corpus.txt', force=True)


def test_lda_train():
    corpus_fp = 'data/1w-data-corpus.txt'
    for topic_num in range(5, 100, 5):
        print('topic num: %d' % topic_num)
        train(corpus_fp, 'model/1w-news-%d-topics-lda.model' % topic_num, topic_num)


def test_infer_doc_topics_mallet():
    doc = """
国金证券：下调中海油服盈利预测，维持“增持”评级

09:11【国金证券：下调中海油服盈利预测，维持“增持”评级】国金证券点评中海油服年报时指出，
低油价影响行业景气度，我们下调盈利预测，预计公司2020-2022年净利润为34.7/44.7/55.9亿元，
下调幅度为16%/11%/14%，同比增长39%/29%/25%；EPS为0.7/0.9/1.2元；
对应PE 17/13/11倍。给予2020年23倍PE，6-12月目标价17元，目标市值810亿元；维持“增持”。
    """
    model_fp = 'model/kuaixun-mallet/2w-kuaixun-20-topics.model'
    lda_model = get_lda_model(model_fp)
    vector = infer_doc_topics(doc, lda_model)
    for topic_id, prob in vector:
        print(prob)
        show_lda_topic(lda_model, topic_id)


def test_infer_doc_topics_gensim():
    doc = """
国金证券：下调中海油服盈利预测，维持“增持”评级

09:11【国金证券：下调中海油服盈利预测，维持“增持”评级】国金证券点评中海油服年报时指出，
低油价影响行业景气度，我们下调盈利预测，预计公司2020-2022年净利润为34.7/44.7/55.9亿元，
下调幅度为16%/11%/14%，同比增长39%/29%/25%；EPS为0.7/0.9/1.2元；
对应PE 17/13/11倍。给予2020年23倍PE，6-12月目标价17元，目标市值810亿元；维持“增持”。
    """
    model_fp = 'model/3M/1000-topics/3M-news-1000-topics-gensim-lda.model'
    print('Loading gensim lda model..')
    lda_model = get_lda_model(model_fp)
    print('Inferring doc topics..')
    start = perf_counter()
    vector = infer_doc_topics(doc, lda_model, False)
    print('Consume: %fs' % (perf_counter() - start))
    print(len(vector))
    for topic_id, prob in vector:
        print(prob)
        show_lda_topic(lda_model, topic_id)


def test_compute_coherence_value():
    corpus_fp = 'data/1w-data-corpus.txt'
    for topic_num in range(5, 100, 5):
        compute_coherence_value('model/1w-news-%d-topics-lda.model' % topic_num, corpus_fp)


def test_split_file():
    fp = './data/1k-data.txt'
    files = split_file(fp, each_count=100)
    assert len(files) == 10


def test_infer_doc_topics_by_glda():
    doc = """
国金证券：下调中海油服盈利预测，维持“增持”评级

09:11【国金证券：下调中海油服盈利预测，维持“增持”评级】国金证券点评中海油服年报时指出，
低油价影响行业景气度，我们下调盈利预测，预计公司2020-2022年净利润为34.7/44.7/55.9亿元，
下调幅度为16%/11%/14%，同比增长39%/29%/25%；EPS为0.7/0.9/1.2元；
对应PE 17/13/11倍。给予2020年23倍PE，6-12月目标价17元，目标市值810亿元；维持“增持”。
    """
    # model_dir, model_config = 'model/1M/glda-128-topics', '1M-news-128-topics-lda.model.conf'
    model_fp = 'model/1M/glda-128-topics.model'
    inf_engine = get_glda_infer_engine(model_fp)
    topic_dis = infer_doc_topics_by_glda(doc, inf_engine)
    print(topic_dis)


def test_read_glda_phi_file():
    phi_fp = 'model/1M/glda-256-topics/256-topics.model.phi'
    phis = read_glda_phi_file(phi_fp)
    for phi in phis:
        print(sum(phi))


def test_read_glda_theta_file():
    theta_fp = 'model/1M/glda-256-topics/256-topics.model.theta'
    for thetas in read_glda_theta_file(theta_fp):
        print(sum(thetas))


def test_get_topic_word_matrix():
    gensim_lda_model_fp = 'model/3M/3M-news-75-topics-lda.model'
    model = get_lda_model(gensim_lda_model_fp)
    matx = get_topic_word_matrix(model)
    for row in matx:
        print(sum(row))


def test_get_word_topic_matrix():
    gensim_lda_model_fp = 'model/3M/3M-news-75-topics-lda.model'
    model = get_lda_model(gensim_lda_model_fp)
    matx = get_word_topic_matrix(model)
    for word_id, row in enumerate(matx):
        if sum(row) > 0.1:
            print(model.id2word[word_id])
            print(sum(row))


def test_calculate_glda_mapv():
    glda_model_fp = 'model/1M/glda-256-topics/256-topics.model'
    mapv = calculate_glda_mapv(glda_model_fp)
    print(mapv)


def test_convert_gensim_lda2familia():
    convert_gensim_lda2familia('model/3M/1000-topics/3M-news-1000-topics-gensim-lda.model')


def test_infer_doc_topics():
    doc = """
国金证券：下调中海油服盈利预测，维持“增持”评级

09:11【国金证券：下调中海油服盈利预测，维持“增持”评级】国金证券点评中海油服年报时指出，
低油价影响行业景气度，我们下调盈利预测，预计公司2020-2022年净利润为34.7/44.7/55.9亿元，
下调幅度为16%/11%/14%，同比增长39%/29%/25%；EPS为0.7/0.9/1.2元；
对应PE 17/13/11倍。给予2020年23倍PE，6-12月目标价17元，目标市值810亿元；维持“增持”。
    """
    model_fp = 'model/3M/75-topics/3M-news-75-topics-lda.model'
    print('Loading gensim lda model..')
    lda_model = get_lda_model(model_fp)
    print('Inferring doc topics..')
    start = perf_counter()
    vector = infer_doc_topics(doc, lda_model, False)
    print('Consume: %fs' % (perf_counter() - start))
    v1 = convert_sparse2dense(vector, 75)
    inf_engine = get_glda_infer_engine(model_fp)
    start = perf_counter()
    topic_dis = infer_doc_topics_by_glda(doc, inf_engine)
    print('Consume: %fs' % (perf_counter() - start))
    v2 = convert_sparse2dense(topic_dis, 75)
    for e1, e2 in zip(v1, v2):
        print('%.4f,%.4f' % (e1, e2))


if __name__ == '__main__':
    # test_batch_get_news_by_oss()
    # test_prepare_corpus()
    # test_lda_train()
    # test_infer_doc_topics()
    # test_compute_coherence_value()
    # test_split_file()
    # test_infer_doc_topics_by_glda()
    # test_read_glda_phi_file()
    # test_read_glda_theta_file()
    # test_get_topic_word_matrix()
    # test_get_word_topic_matrix()
    # test_calculate_glda_mapv()
    # test_get_news_by_url()
    # test_infer_doc_topics_gensim()
    # test_convert_gensim_lda2familia()
    test_infer_doc_topics()
    os._exit(0)
